from influxdb_client import InfluxDBClient
from influxdb_client.client.write_api import SYNCHRONOUS
import csv
import argparse
from datetime import datetime

class InfluxDBLoader:
    def __init__(self, url: str, token: str, org: str):
        """Initialize InfluxDB connection parameters"""
        self.url = url
        self.token = token
        self.org = org
        self.client = None
        self.query_api = None

    def connect(self):
        """Establish connection to InfluxDB"""
        try:
            self.client = InfluxDBClient(url=self.url, token=self.token, org=self.org)
            self.query_api = self.client.query_api()
            return True
        except Exception as e:
            print(f"Failed to connect to InfluxDB: {e}")
            return False

    def query_and_export(self, bucket: str, measurement: str, device_id:str, start_time: str, stop_time: str):
        """
        Query data from InfluxDB and export to CSV
        
        Args:
            bucket: Database/bucket name
            measurement: Table/measurement name
            start_time: Start time in RFC3339 format (e.g., '2023-01-01T00:00:00Z')
            stop_time: Stop time in RFC3339 format (e.g., '2023-01-02T00:00:00Z')
        """
        if not self.client:
            print("Not connected to InfluxDB. Call connect() first.")
            return False

        try:
            # Build Flux query
            query = f'''
            from(bucket: "{bucket}")
              |> range(start: {start_time}, stop: {stop_time})
              |> filter(fn: (r) => r._measurement == "{measurement}")
              |> filter(fn: (r) => r["_field"] == "acc_x" or r["_field"] == "acc_y" or r["_field"] == "acc_z" or r["_field"] == "titl_x" or r["_field"] == "titl_y" or r["_field"] == "titl_z" or r["_field"] == "counter")
              |> filter(fn: (r) => r["device_id"] == "{device_id}")
              |> pivot(rowKey: ["_time"], columnKey: ["_field"], valueColumn: "_value")
            '''

            # Execute query
            result = self.query_api.query(query=query, org=self.org)

            # Generate filename from measurement and date
            date_str = datetime.now().strftime("%Y%m%d")

            stime_str=datetime.fromisoformat(start_time)
            stime_output=stime_str.strftime("%Y%m%d%H%M%S")
            etime_str=datetime.fromisoformat(stop_time)
            etime_output=etime_str.strftime("%Y%m%d%H%M%S")

            filename = f"angle-scl3300d01-{measurement.lower()}-{stime_output}_{etime_output}.csv"

            # Write results to CSV
            with open('../data_'+date_str+'/'+filename, 'w', newline='') as csvfile:
                writer = csv.writer(csvfile)
                
                # Write header
                writer.writerow(['_receviced','acc_x', 'acc_y', 'acc_z', 'angle_x', 'angle_y', 'angle_z','counter','devicd','magnitude','ts'])
                
                # Write data
                for table in result:
                    for record in table.records:
                        writer.writerow([
                            '',
                            record.values['acc_x'],
                            record.values['acc_y'],
                            record.values['acc_z'],
                            record.values['titl_x'],
                            record.values['titl_y'],
                            record.values['titl_z'],
                            record.values['counter'],
                            record.values['device_id'],
                            '',
                            int(record.get_time().timestamp() * 1000)
                        ])

            print(f"Data successfully exported to {filename}")
            return True

        except Exception as e:
            print(f"Error querying/exporting data: {e}")
            return False

    def close(self):
        """Close the InfluxDB connection"""
        if self.client:
            self.client.close()
            self.client = None
            self.query_api = None

def data_with_influxdb(bucket:str,measurement:str,deviceid:str,stime:str,etime:str):
    """
    get data from InfluxDB.
    
    Args:
        bucket: influxdb bucket's name
        measurement: influxdb measurement'name
        stime: data's starttime
        etime: data's endtime
    """

    URL = "http://192.168.195.169:8086"
    TOKEN = "tM4zhHUjjGomPm0QnlUgLO7Nb9kd5pjwSE1R3n9TYl__PMsEP5cDV-OqJzfPLJDvmQndpyq9Nwx_rtnkSREfvg=="
    ORG = "SmartDesk"

    # Create loader instance
    loader = InfluxDBLoader(url=URL, token=TOKEN, org=ORG)

    # Connect to InfluxDB
    if loader.connect():
        try:
            # Query and export data
            loader.query_and_export(
                bucket=bucket,
                measurement=measurement,
                device_id=deviceid,
                start_time=stime,
                stop_time=etime
            )
        finally:
            # Ensure connection is closed
            loader.close()

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="从influxdb获取数据")
    parser.add_argument('--bucket',type=str,required=True,help='库名称')
    parser.add_argument('--measurement',type=str,required=True,help='表名称')
    parser.add_argument('--deviceid',type=str,required=True,help='设备Mac地址')
    parser.add_argument('--stime',type=str,required=True,help='数据开始时间')
    parser.add_argument('--etime',type=str,required=True,help='数据结束时间')

    args = parser.parse_args()
    #data_with_influxdb(bucket=args.bucket,measurement=args.measurement,stime=args.stime,etime=args.etime)